'''
Created on May 14, 2009

@author: mkiyer
'''
import re
import sys

class ParseError(Exception):
    """Exception raised for errors parsing input
    
    Attributes: message -- explanation of the error
    """
    def __init__(self, message):
        self.message = message
    def __str__(self):
        return repr(self.message)

class ElandMapping(object):
    '''Describes the location of a sequence alignment
    
    Attributes:
    chrom - chromosome name of the alignment in string format 'chrN'
    chrom_start - integer location of start of alignment based on (+) strand
    chrom_end - integer location of end of alignment
    num_mismatches - number of alignment mismatches
    '''
    def __init__(self, chrom, chrom_start, chrom_end, strand, num_mismatches=0):
        self.chrom = chrom
        self.chrom_start= chrom_start
        self.chrom_end = chrom_end
        self.strand = strand
        self.num_mismatches = num_mismatches
    
    def __repr__(self):
        return "ElandAlignment('%s', '%s', '%s', '%s', '%s')" % \
            (self.chrom, self.chrom_start, self.chrom_end, self.strand, 
             self.num_mismatches)
    
    def maps_to_genome(self):
        return self.chrom.startswith('hs_ref')

class ElandRead(object):
    def is_mapping(self):
        pass

class ElandExportRead2(ElandRead):    
    def filter_quality_filter(self):
        return self.passes_filter
    def filter_maps_to_genome(self):
        return self.chrom.startswith('hs_ref')
    def filter_maps_to_genes(self):
        return self.chrom.startswith('ucsc')

class ElandMultiRead(ElandRead):
    def get_match_counts(self):
        self.match_code = map(int, self.match_code.split(':'))

class ElandParser2(object):
    def __init__(self):
        self.filters = []

    def add_filter(self, filter_func):
        self.filters.append(filter_func)

    def parse_file(self, thisfhd):
        i = 0
        m = 0
        for thisline in thisfhd:
            # count lines for status updates
            i+=1
            if i == 100000:
                m += 1
                sys.stderr.write(" %d\n" % (m*100000))
                i=0
            # process read
            read = self.parse_line(thisline)
            passed_filters = True
            for filter_func in self.filters:
                if not filter_func(read):
                    passed_filters = False
                    break
            if passed_filters:
                yield read

    def configure(self, desired_fields):
        '''
        desired_fields: list of field names to capture
        '''
        self.desired_fields = set([])
        for i, field_name in enumerate(self._field_names):
            if field_name in desired_fields:
                self.desired_fields.add(i)
            else:
                raise NameError('field ' + field_name + 
                                ' does not exist in ' + 
                                str(self._field_names)) 
        
    def parse_line(self, thisline):
        if not thisline:
            raise ParseError("line is None")
        thisline = thisline.strip()
        if not thisline:
            raise ParseError("blank line")
        if thisline.startswith("#"):
            return None

        read = ElandRead()
        # populate attributes of the read
        thisfields = thisline.split('\t')
        for i in self.desired_fields:
            if i >= len(thisfields):
                value = None
            else:
                value = thisfields[i]
            setattr(read, self._field_names[i], thisfields[i])


class ElandExportParser(ElandParser2):
    
    _field_names = ['machine',
                    'run_number',
                    'lane',
                    'tile',
                     'cluster_x_coord',
                     'cluster_y_coord',
                     'index_string',
                     'read_number',
                     'seq',
                     'quality_string',
                     'chrom',
                     'contig',
                     'chrom_start',
                     'strand',
                     'match_descriptor',
                     'single_read_alignment_score',
                     'paired_read_alignment_score',
                     'partner_chrom',
                     'partner_contig',
                     'partner_offset',
                     'partner_strand',
                     'passes_filter']

    def __init__(self):
        super(ElandExportParser, self).__init__()

    def get_num_mismatches(self, match_descriptor):
        num_mismatches = 0
        for c in match_descriptor:
            if c in frozenset(['A', 'T', 'G', 'C']):
                num_mismatches += 1
        return num_mismatches

    def get_mappings(self, thisline):
        thisline = thisline.strip()
        fields = thisline.split('\t')
        seqlen = len(fields[8])
        chrom = fields[10]
        chrom_start = (int(fields[12])-1) if fields[12] else 0
        strand = fields[13]
        num_mismatches = self.get_num_mismatches(fields[14])
        return ElandMapping(chrom=chrom,
                            chrom_start=chrom_start,
                            chrom_end=chrom_start + seqlen,
                            strand=strand,
                            num_mismatches=num_mismatches)

    def parse_line(self, thisline):
        thisline = thisline.strip()
        fields = thisline.split('\t')
        read = ElandExportRead2()
        read.machine = fields[0]
        read.run_number = fields[1]
        read.lane = fields[2]
        read.tile = fields[3]
        read.cluster_x_coord = fields[4]
        read.cluster_y_coord = fields[5]
        read.index_string = fields[6]
        read.read_number = fields[7]
        read.seq = fields[8]
        read.quality_string = fields[9]
        read.chrom = fields[10]
        read.contig = fields[11]
        # eland format is one-based so convert to zero-based
        read.chrom_start = (int(fields[12])-1) if fields[12] else 0        
        read.strand = fields[13]
        read.match_descriptor = fields[14]
        read.single_read_alignment_score = fields[15]
        read.paired_read_alignment_score = fields[16]
        read.partner_chrom = fields[17]
        read.partner_contig = fields[18]
        read.partner_offset = fields[19]
        read.partner_strand = fields[20]
        read.passes_filter = (fields[21] == 'Y')
        # construct an alignment for this read
        num_mismatches = self.get_num_mismatches(read.match_descriptor)
        read.matches = [ElandMapping(read.chrom, read.chrom_start,
                                     read.chrom_start + len(read.seq),
                                     read.strand, num_mismatches)]
        return read

class ElandMultiParser(ElandParser2):

    _field_names = ['name',
                    'seq',
                    'match_code',
                    'matches']
    
    def __init__(self):
        super(ElandMultiParser, self).__init__()
        self.do_stats = False
        self.desired_fields = self._field_names
        
    def parse_line(self, thisline):
        '''
        >PATHBIO-SOLEXA1:4:1:10:1986#0/1        CAGAGAGAGAGAGACAGAGAGACAGAGAGAGAGAGACAGA        7:127:255       hs_ref_chr1.fa:3635418F1,10682680F1,31818384R1,49972546F1,50314077
        F1,58433390F1,156724394R1,179627675R1,214013468R1,hs_ref_chr10.fa:6549303F1,45731419R1,50763139R1,50984872R1,51349836F1,86186898F1,93358899R1,123793042R1,hs_ref_chr11.fa:
        1605111F1,26027502R1,43609718R1,115506815R1,131265251R1,hs_ref_chr12.fa:1341866R1,68248605R1,85294713F1,97986612F1,109309485F1,109980722F1,117680020R1,123942466R1,1314637
        10F1,hs_ref_chr13.fa:25364599R1,93423164R1,hs_ref_chr14.fa:32284445R1,70163306F1,83934289F1,87751841F1,hs_ref_chr15.fa:21350544R1,57929214F1,69538722R1,84241244R1,hs_ref_
        chr16.fa:16440972F1,20232067R1,24528368F1,62384015R1,76162293R1,84410611R1,hs_ref_chr17.fa:6326101R1,11171972R1,17308823R1,62516266F1,hs_ref_chr18.fa:18046154F0,27060570F
        1,28760280R1,47512990F1,hs_ref_chr19.fa:15654880F1,42473947F1,42478998F1,59103884R0,hs_ref_chr2.fa:44952099R1,78595375R1,78651084F1,78651114F1,125338576R1,155219795R1,163
        973179F1,184702065F1,200317413F1,202635216R1,207996572F1,214708675R1,222661222R1,222905782F1,228216330F0,hs_ref_chr20.fa:15355652R1,40305143F1,hs_ref_chr21.fa:24365316R1,
        46225422R1,hs_ref_chr22.fa:42849719F1,hs_ref_chr3.fa:81428319R1,107762952F1,117011765R1,119887661F1,hs_ref_chr4.fa:2215404F1,26711287F0,69987368R1,80493532F1,82760322F0,1
        56532726F1,181118735R1,hs_ref_chr5.fa:4867371R1,51095771F1,96976337F1,161201517R1,161676256R1,hs_ref_chr6.fa:9579077F1,55814659F1,91180692F1,115387109R1,124281822R1,hs_re
        f_chr7.fa:1597962R1,22593180R1,29972278F1,41327262F1,49168874R1,67231050F1,70464708F1,91088100F1,129675602R1,156946722R1,156946754R0,156946996R1,hs_ref_chr8.fa:2985261R1,
        10887175R1,101919154R1,hs_ref_chr9.fa:7261673F1,86097201R1,hs_ref_chrX.fa:1180448F1,1180962F1,27994328R1,37551058R0,62625745F1,67847833F1,118686303F1,124312775F1,12460464
        6R1,136242290R1,141295855F1,151520642F1,hs_ref_chrY.fa:1180448F1,1180962F1,10634241R1,12885845F1,humRibosomal.fa:15469R1
        '''
        thisline = thisline.strip()
        fields = thisline.split('\t')
        
        read = ElandMultiRead()
        read.name = fields[0]
        read.seq = fields[1]
        read.match_code = fields[2]
        read.matches = []

        if (read.match_code == 'NM' or read.match_code == 'QC' or 
            read.match_code == 'RM' or (len(fields) < 4)):
            return read

        chrom = None
        seqlen = len(read.seq)
      
        for matchstring in fields[3].split(','):
            # check if contig (chromosome) is present in this string,
            # and if so suck it out of the matchstring
            matchstring = matchstring.split(':')
            if len(matchstring) == 2:
                chrom, matchstring = matchstring
            else:
                matchstring = matchstring[0]
            # separate strand and number of mismatches from position
            # fortunately the number of mismatches is guaranteed to be
            # 0, 1, or 2, so it will also consume just one character
            strand = matchstring[-2]
            chrom_start = int(matchstring[:-2])-1
            chrom_end = chrom_start + seqlen
            num_mismatches = matchstring[-1]
            # append a new alignment to the list of matches
            read.matches.append(ElandMapping(chrom, chrom_start, 
                                             chrom_start + seqlen,
                                             strand, num_mismatches))
        return read



class ElandParser(object):
    '''
    Parser for eland file formats
    '''
    def __init__(self):
        '''
        Constructor
        '''
        pass
    
    @staticmethod
    def parseElandExportFile(fhd):
        '''
        generator function that returns list of reads
        '''
        for line in fhd:
            read = ElandExportRead()
            read.parseLine(line)
            yield read

    @staticmethod
    def getElandExportFileStats(fhd, outfhd):
        '''
        report statistics on mapping/non-mapping/filter/non-filter reads
        
        fhd: file handle descriptor to eland export file
        outfhd: file handle descriptor to output results
        '''
        total = 0
        unique = set([])
        pf = 0
        pf_nm_contam = 0
        pf_nm_ribosome = 0
        pf_nm_mito = 0
        pf_nm = 0
        pf_nm_qc = 0
        pf_nm_multimap = 0
        pf_dna = 0
        pf_genes = 0
        pf_chrom = {}
        pf_dna_unique = set([])
        fails_filter_but_maps = 0

        # parse file
        for read in ElandParser.parseElandExportFile(fhd):
            total += 1
            unique.add(read.read_seq)
            # only process reads passing filter
            if read.passes_filter:
                pf += 1
                # maps to contaminant file
                if read.match_chromosome == 'newcontam.fa':
                    pf_nm += 1
                    pf_nm_contam += 1
                    continue
                # maps to ribosomal
                if read.match_chromosome == 'humRibosomal.fa':
                    pf_nm += 1
                    pf_nm_ribosome += 1
                    continue
                # maps to mitochondrial
                if read.match_chromosome == 'chrM.fa':
                    pf_nm += 1
                    pf_nm_mito += 1
                    continue
                # TODO: reads matching one of the adaptors?
                # has QC issues
                if read.match_chromosome == 'QC':
                    pf_nm += 1
                    pf_nm_qc += 1
                    continue
                # non-mapping reads
                if read.match_chromosome == 'NM':
                    pf_nm += 1
                    continue
                # reads with too many multi-maps (e.g. 12:200:255)
                if re.match(r"^[0-9]", read.match_chromosome):
                    pf_nm += 1
                    pf_nm_multimap += 1
                    continue
                # TODO: handle reads with match_chromosome == allMrna.fa
                # reads mapping to genome
                if (read.match_chromosome.startswith('hs_ref')):
                    pf_dna += 1
                    pf_dna_unique.add(read.read_seq)
                    chrom_name = re.match(r'hs_ref_(\w+).fa', read.match_chromosome).group(1)
                    if not chrom_name in pf_chrom:
                        pf_chrom[chrom_name] = 0
                    pf_chrom[chrom_name] += 1
                    continue
                # reads mapping to ucsc genes
                if (read.match_chromosome.startswith('ucsc')):
                    pf_genes += 1
                    continue
            else:
                if (read.match_chromosome.startswith('hs_ref')):
                    fails_filter_but_maps += 1
                    continue
        
        # output statistics
        percent = lambda x,y: round(float(100*x)/y, 2)
        print >>outfhd, 'Total reads:', total
        print >>outfhd, 'unique reads:', len(unique), '(', percent(len(unique), total), '%)'
        print >>outfhd, 'PF reads:', pf, '(', percent(pf, total), '%)'
        print >>outfhd, 'PF reads non-mapping:', pf_nm, '(', percent(pf_nm, total), '%)'
        print >>outfhd, 'PF reads non-mapping QC:', pf_nm_qc, '(', percent(pf_nm_qc, total), '%)'
        print >>outfhd, 'PF reads non-mapping multimap:', pf_nm_multimap, '(', percent(pf_nm_multimap, total), '%)'
        print >>outfhd, 'PF reads non-mapping contam:', pf_nm_contam, '(', percent(pf_nm_contam, total), '%)'
        print >>outfhd, 'PF reads non-mapping ribosome:', pf_nm_ribosome, '(', percent(pf_nm_ribosome, total), '%)'
        print >>outfhd, 'PF reads non-mapping mitochondria:', pf_nm_mito, '(', percent(pf_nm_mito, total), '%)'
        print >>outfhd, 'PF reads mapping:', pf-pf_nm, '(', percent(pf-pf_nm, total), '%)'
        print >>outfhd, 'PF reads mapping dna:', pf_dna, '(', percent(pf_dna, total), '%)'
        print >>outfhd, 'PF reads mapping genes:', pf_genes, '(', percent(pf_genes, total), '%)'        
        print >>outfhd, 'PF reads mapping dna and unique:', len(pf_dna_unique), '(', percent(len(pf_dna_unique), pf_dna), '% of dna mapping reads)'
        print >>outfhd, 'Fails filter but maps dna:', fails_filter_but_maps, '(', percent(fails_filter_but_maps, total), '%)'
        for chrom, count in pf_chrom.iteritems():
            print >>outfhd, 'PF reads mapping', chrom, ':', count, '(', percent(count, pf_dna), '% of dna mapping reads)'

    @staticmethod
    def parseElandResultFile(fhd):
        '''
        generator func that returns list of reads
        '''
        for thisline in fhd:
            if not thisline:
                continue
            if thisline.startswith("#") or thisline.startswith("track") or thisline.startswith("browser"):
                # TODO: comment line
                continue            
            thisline = thisline.rstrip()
            if not thisline:
                # blank line
                continue
            read = ElandResultRead()
            read.parseLine(thisline)
            yield read
    

class ElandResultRead(object):
    '''
    stores information about a single read from the eland result file
    '''
    def __init__(self):
        pass
    
    def hasUniqueBestMatch(self):
        if self.matchCode == 'U0' or self.matchCode == 'U1' or self.matchCode == 'U2':
            return True
        return False
    
    def hasMismatches(self):
        if self.matchCode == 'U1' or self.matchCode == 'U2':
            return True
        return False
    
    def parseLine(self, thisline):
        '''
        >PATHBIO-SOLEXA_20E6WAAXX:8:1:914:878   GGCTGATAGTGTCAGATTTCGTGTTGATTTTTTTTT    U2      0       0       1       hs_ref_chr6.fa  44343560        R       ..      8C      30T

        Fields:
        0. Sequence name (derived from file name and line number if format is not fasta) 
        1. Sequence 
        2. Type of match codes:
            NM-No match found
            QC-No matching done: QC failure (too many Ns) 
            RM-No matching done: repeat masked (may be seen if repeatFile.txt was specified) 
            U0-Best match found was a unique exact match 
            U1-Best match found was a unique 1-error match 
            U2-Best match found was a unique 2-error match 
            R0-Multiple exact matches found 
            R1-Multiple 1-error matches found, no exact matches 
            R2-Multiple 2-error matches found, no exact or 1-error matches 
        3. Number of exact matches found 
        4. Number of 1-error matches found 
        5. Number of 2-error matches found 
        The following fields are only used if a unique best match was found: 
        6. Genome file in which match was found 
        7. Position of match (bases in file are numbered starting at 1) 
        8.Direction of match (F=forward strand, R=reverse) 
        9.How N characters in read were interpreted ("."=not applicable, "D"=Detection, "I"=Insertion) 
        The following field is only used in the case of a unique inexact match: 
        10.Position and type of first substitution error (A numeral refers to a run of matching 
        bases, an upper case base or N refers to a base in the reference that differs from the 
        read. For example, 11A: after 11 matching bases, base 12 is A in the reference but 
        not in the read)
        '''
        fields = thisline.split('\t')
        self.name = fields[0]
        self.seq = fields[1]
        self.matchCode = fields[2]
        if self.matchCode == 'QC':
            return
        self.numMatches = [int(fields[3]), int(fields[4]), int(fields[5])]
        if self.hasUniqueBestMatch():
            self.chrom = fields[6]
            self.position = int(fields[7]) - 1            
            self.strand = fields[8]
            self.uncalledBaseHandling = fields[9]
            if self.hasMismatches():
                self.mismatches = fields[10]

class ElandExportRead(object):
    '''
    stores information about a single read from the eland export file
    '''
    def __init__(self):
        '''
        Constructor
        '''
        pass

    def parseLine(self, line):
        '''
        Parse a line from the eland export file and populate the class data
        
        File format is tab-delimited fields as specified in the Solexa Pipeline version 1.1 documentation:
        0. machine
        1. run_number
        2. lane
        3. tile
        4. x coordinate of cluster
        5. y coordinate of cluster
        6. index string (bland for a non-indexed run)
        7. read number (1 or 2 for paired-read analysis, blank for a single-read analysis)
        8. read sequence
        9. quality string
        10. match chromosome (or code describing why no match occurred)
        11. match contig
        12. match position (always with respect to forward strand, eland numbering starts at one)
        13. match strand ('F' for forward, 'R' for reverse)
        14. match descriptor
        15. single-read alignment score
        16. paired-read alignment score
        17. partner chromosome
        18. partner contig
        19. partner offset
        20. partner strand
        22. filtering ('Y' for yes, 'N' for no)        
        '''
        fields = line.strip().split('\t')
        self.machine = fields[0]
        self.run_number = fields[1]
        self.lane = fields[2]
        self.tile = fields[3]
        self.clusterxcoord = fields[4]
        self.clusterycoord = fields[5]
        self.index_string = fields[6]
        self.read_number = fields[7]
        self.read_seq = fields[8]
        self.quality_string = fields[9]
        self.match_chromosome = fields[10]
        self.match_contig = fields[11]
        # eland format is one-based so convert to zero-based
        self.match_position = (int(fields[12])-1) if fields[12] else 0        
        self.match_strand = fields[13]
        self.match_descriptor = fields[14]
        self.single_read_alignment_score = fields[15]
        self.paired_read_alignment_score = fields[16]
        self.partner_chromosome = fields[17]
        self.partner_contig = fields[18]
        self.partner_offset = fields[19]
        self.partner_strand = fields[20]
        self.passes_filter = (fields[21] == 'Y')